use anyhow::Result;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine;
use chrono::{DateTime, Local, TimeZone};
use serde::{de, ser, Deserialize, Serialize};
use serde_with::DisplayFromStr;

use chacha20poly1305::{
    aead::{Aead, AeadCore, KeyInit, OsRng},
    ChaCha20Poly1305,
};
use serde_with::serde_as;

const KEY: &[u8] = b"01234567890123456789012345678901";

#[serde_as]
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct User {
    #[serde(rename = "userName")] // 字段重命名
    name: String,

    #[serde(default)] // 反序列化时提供默认值
    age: u8,

    #[serde(serialize_with = "serialize", deserialize_with = "deserialize")]
    dob: DateTime<Local>,
    skills: Vec<String>,

    state: MyState,

    #[serde(serialize_with = "b64_encode", deserialize_with = "b64_decode")]
    data: Vec<u8>,

    #[serde(
        serialize_with = "serialize_encrypt",
        deserialize_with = "serialize_decrypt"
    )]
    sensitive: String,

    #[serde_as(as = "DisplayFromStr")]
    url: http::Uri,
}

// 自定义序列化函数
fn serialize<S>(dt: &DateTime<Local>, serializer: S) -> Result<S::Ok, S::Error>
where
    S: ser::Serializer,
{
    const FORMAT: &str = "%Y-%m-%d %H:%M:%S";
    serializer.serialize_str(&dt.format(FORMAT).to_string())
}

fn deserialize<'de, D>(deserializer: D) -> Result<DateTime<Local>, D::Error>
where
    D: de::Deserializer<'de>,
{
    use chrono::{DateTime, Local, NaiveDateTime};

    let s = String::deserialize(deserializer)?;

    // 尝试解析为带有时区的字符串
    if let Ok(dt) =
        DateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S%.f %z").map(|dt| dt.with_timezone(&Local))
    {
        return Ok(dt);
    }

    if let Ok(dt) =
        DateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S %z").map(|dt| dt.with_timezone(&Local))
    {
        return Ok(dt);
    }

    // 尝试解析为无时区的字符串，并假设为本地时间
    if let Ok(dt) = NaiveDateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S%.f") {
        return match Local.from_local_datetime(&dt).single() {
            Some(dt) => Ok(dt),
            None => Err(de::Error::custom(format!("Ambiguous time: {s}"))),
        };
    }

    if let Ok(dt) = NaiveDateTime::parse_from_str(&s, "%Y-%m-%d %H:%M:%S") {
        return match Local.from_local_datetime(&dt).single() {
            Some(dt) => Ok(dt),
            None => Err(de::Error::custom(format!("Ambiguous time: {s}"))),
        };
    }

    Err(de::Error::custom(format!(
        "Failed to parse '{s}' as local datetime"
    )))
}

//encrypt with chacha20poly1305
fn encrypt(data: &[u8]) -> Result<String> {
    let cipher = ChaCha20Poly1305::new(KEY.into());
    let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
    let ciphertext = cipher.encrypt(&nonce, data).unwrap();
    let nonce_cypertext: Vec<_> = nonce.iter().copied().chain(ciphertext).collect();
    let encoded = URL_SAFE_NO_PAD.encode(nonce_cypertext);
    Ok(encoded)
}

fn decrypt(encoded: &str) -> Result<Vec<u8>> {
    let decoded = URL_SAFE_NO_PAD.decode(encoded.as_bytes())?;
    let cipher = ChaCha20Poly1305::new(KEY.into());
    let nonce = decoded[..12].into();
    let decrypted = cipher.decrypt(nonce, &decoded[12..]).unwrap();
    Ok(decrypted)
}

fn serialize_encrypt<S>(data: &String, serializer: S) -> Result<S::Ok, S::Error>
where
    S: ser::Serializer,
{
    let encoder = encrypt(data.as_bytes()).unwrap();
    serializer.serialize_str(&encoder)
}

fn serialize_decrypt<'de, D>(deserializer: D) -> Result<String, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let encrypted = String::deserialize(deserializer)?;
    let decrypted = decrypt(&encrypted).map_err(de::Error::custom)?;
    let decrypted = String::from_utf8(decrypted).map_err(de::Error::custom)?;
    Ok(decrypted)
}

fn b64_encode<S>(data: &Vec<u8>, serializer: S) -> Result<S::Ok, S::Error>
where
    S: ser::Serializer,
{
    let encoder = URL_SAFE_NO_PAD.encode(data);
    serializer.serialize_str(&encoder)
}

fn b64_decode<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let encoded = String::deserialize(deserializer)?;
    let decoded = URL_SAFE_NO_PAD
        .decode(encoded.as_bytes())
        .map_err(serde::de::Error::custom)?;
    Ok(decoded)
}

#[derive(Serialize, Deserialize, Debug, PartialEq)]
#[serde(rename_all = "snake_case")]
enum MyState {
    // #[serde(rename="init")]
    Init(String),
    #[serde(rename = "running")]
    Running(Vec<String>),
    // #[serde(rename="done")]
    Done(u32),
}

fn main() -> Result<()> {
    let user = User {
        name: "abc".to_string(),
        age: 10,
        dob: Local::now(),
        skills: vec!["Rust".to_string(), "Python".to_string()],
        state: MyState::Init("init".to_string()),
        data: vec![1, 2, 3],
        sensitive: "dgj".to_string(),
        url: "http://www.baidu.com".parse()?,
    };
    let json = serde_json::to_string(&user)?;
    println!("{json}");
    let user1: User = serde_json::from_str(&json)?;
    println!("{user1:?}");
    // assert_eq!(user1, user);

    Ok(())
}
